from typing import Sequence, Union

import numpy as np
import sapien.core as sapien
from sapien.core import Pose
from transforms3d.euler import euler2quat

from mani_skill2.agents.base_agent import BaseAgent
from mani_skill2.agents.configs.cocube import defaults


def get_entities_by_names(
    entities: Sequence[sapien.Entity], names: Union[str, Sequence[str]]
):
    assert isinstance(entities, (list, tuple)), type(entities)
    if isinstance(names, str):
        names = [names]
        flag = True  # singular
    else:
        flag = False
    ret = [None for _ in names]

    for entity in entities:
        name = entity.get_name()
        if name in names:
            ret[names.index(name)] = entity
    return ret[0] if flag else ret





class Cocube(BaseAgent):
    _config: defaults.CocubeDefaultConfig

    def __init__(
        self, scene, control_freq, control_mode=None, fix_root_link=False, config=None
    ):
        if control_mode is None:  # if user did not specify a control_mode
            control_mode = "pd_joint_vel"
        super().__init__(
            scene,
            control_freq,
            control_mode=control_mode,
            fix_root_link=fix_root_link,
            config=config,
        )

    @classmethod
    def get_default_config(cls):
        return defaults.CocubeDefaultConfig()

    def _after_init(self):
        super()._after_init()

